import subprocess
import re
import logging
from typing import List, Union, Optional

class ADBUtils:
    """ADB操作工具类"""
    
    def __init__(self):
        self.log = logging.getLogger("ADBUtils")
    
    @staticmethod
    def sanitize_device_id(device_id: str) -> str:
        """清理设备ID中的非法字符"""
        return re.sub(r'[\\/*?:"<>|\n\r:]', '_', device_id.strip())
    
    def run_adb(self, command: List[str], device_id: str = "") -> Optional[Union[str, bytes]]:
        """
        执行ADB命令
        
        :param command: ADB命令列表
        :param device_id: 目标设备ID
        :return: 命令输出（文本或二进制）
        """
        try:
            # 构建基础命令
            cmd = ["adb"]
            if device_id:
                cmd.extend(["-s", device_id])
            cmd.extend(command)
            
            # 执行命令并捕获输出
            result = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=15
            )
            
            # 检查返回状态
            if result.returncode != 0:
                error_msg = result.stderr.decode('utf-8', errors='ignore').strip()
                self.log.error(f"ADB执行失败 [{result.returncode}]: {error_msg}")
                return None
                
            # 直接返回原始二进制输出
            return result.stdout
            
        except subprocess.TimeoutExpired as te:
            self.log.error(f"ADB命令执行超时: {str(te)}")
            return None
        except Exception as e:
            self.log.error(f"ADB命令异常: {str(e)}")
            return None
    
    def capture_screen(self, output_path: str, device_id: str = "") -> bool:
        """
        捕获设备屏幕截图
        
        :param output_path: 截图保存路径
        :param device_id: 目标设备ID
        :return: 是否成功
        """
        try:
            # 执行截屏命令
            screenshot_data = self.run_adb(
                ["exec-out", "screencap", "-p"], 
                device_id
            )
            
            if screenshot_data is None:
                self.log.error("截屏失败: 未获取到数据")
                return False
                
            # 确保是二进制数据
            if not isinstance(screenshot_data, bytes):
                self.log.error(f"截屏返回了非二进制数据: {type(screenshot_data)}")
                return False
                
            # 保存截图文件
            with open(output_path, "wb") as f:
                f.write(screenshot_data)
            self.log.info(f"截图保存至: {output_path}")
            return True
            
        except Exception as e:
            self.log.error(f"截屏异常: {str(e)}")
            return False
    
    def wake_screen(self, device_id: str = ""):
        """通过电源键唤醒设备屏幕"""
        try:
            # 使用文本模式执行唤醒命令
            self.run_adb(["shell", "input", "keyevent", "224"], device_id)
            self.log.info(f"设备 {device_id} 屏幕已唤醒")
        except Exception as e:
            self.log.error(f"屏幕唤醒失败: {str(e)}")
    
    def tap_screen(self, x: int, y: int, device_id: str = "") -> bool:
        """
        点击屏幕指定位置
        
        :param x: 横坐标
        :param y: 纵坐标
        :param device_id: 目标设备ID
        :return: 是否成功
        """
        try:
            # 使用文本模式执行点击命令
            self.run_adb(["shell", "input", "tap", str(x), str(y)], device_id)
            self.log.info(f"点击操作: ({x}, {y})")
            return True
        except Exception as e:
            self.log.error(f"点击失败({x},{y}): {str(e)}")
            return False
    
    def get_connected_devices(self) -> List[str]:
        """获取已连接的ADB设备列表"""
        try:
            # 获取文本格式的设备列表
            output = self.run_adb(["devices"])
            
            if output is None:
                self.log.error("获取设备列表失败: 无输出")
                return []
                
            # 确保是字符串类型
            if isinstance(output, bytes):
                output = output.decode('utf-8', errors='ignore')
                
            devices = []
            for line in output.splitlines():
                if "\tdevice" in line:
                    device_id = line.split("\t")[0].strip()
                    if device_id:
                        devices.append(device_id)
                        
            return devices
            
        except Exception as e:
            self.log.error(f"获取设备列表失败: {str(e)}")
            return []

    def go_home(self, device_id: str = "") -> bool:
        """
        按 Home 键返回桌面
        :param device_id: 目标设备ID
        :return: 是否成功
        """
        try:
            # KEYCODE_HOME 键码为 3
            self.run_adb(["shell", "input", "keyevent", "3"], device_id)
            self.log.info(f"按Home键返回桌面: {device_id}")
            return True
        except Exception as e:
            self.log.error(f"按Home键失败: {str(e)}")
            return False